import os
import xml.etree.ElementTree as ET
import json
from typing import Dict, Any, Generator

LOCAL_JSON_BATH_PATH = r"C:\local_pubmed_bath_json"

def load_large_xml_data_to_batch_json(file_path: str, batch_size: int = 100):
    """读取XML文件并批量保存为JSON文件"""
    article_batch = []
    batch_counter = 1
    # 提取文件名作为批次前缀
    xml_filename = os.path.splitext(os.path.basename(file_path))[0]
    try:
        for event, elem in ET.iterparse(file_path, events=('start', 'end')):
            if event == 'end' and elem.tag == 'PubmedArticle':
                parsed_article = parse_pubmed_article(elem)
                pmid = parsed_article.get('PMID')
                year = parsed_article.get('Year')
                title = parsed_article.get('JournalTitle')
                article_title = parsed_article.get('ArticleTitle')
                abstract_text = parsed_article.get('AbstractText')
                try:
                    if (pmid and int(pmid) > 31000000 and title
                            and article_title and year and abstract_text):
                        title_display = article_title[:50] if article_title else 'N/A'
                        print(f"解析文章: PMID={pmid},{year}, {title} ,标题={title_display}...")
                        article_batch.append(parsed_article)
                        if len(article_batch) >= batch_size:
                            save_articles_batch_simple(article_batch, batch_counter, LOCAL_JSON_BATH_PATH,xml_filename)
                            article_batch = []
                            batch_counter += 1
                except (ValueError, TypeError):
                    pass
                elem.clear()
    except ET.ParseError as e:
        print(f"XML解析错误: {e}")
    except FileNotFoundError:
        print(f"文件未找到: {file_path}")
    except MemoryError:
        print("内存不足，无法处理大文件")
    if article_batch:
        save_articles_batch_simple(article_batch, batch_counter, LOCAL_JSON_BATH_PATH,xml_filename)
        print(f"最终批次 {batch_counter} 已保存 {len(article_batch)} 篇文章")

def save_articles_batch_simple(articles: list, batch_id: int, base_path: str, prefix: str):
    """批量保存文章数据到单个JSON文件（只包含指定字段）"""
    if not articles:
        return
    os.makedirs(base_path, exist_ok=True)
    # 生成基础文件名
    # 使用 prefix 生成文件名
    base_file_name = f"{prefix}_{batch_id:04d}.json"
    file_path = os.path.join(base_path, base_file_name)
    simple_articles = []
    for article in articles:
        simple_article = {
            'PMID': article.get('PMID'),
            'Year': article.get('Year'),
            'Title': article.get('JournalTitle'),
            'ArticleTitle': article.get('ArticleTitle'),
            'AbstractText': article.get('AbstractText')
        }
        simple_articles.append(simple_article)
    with open(file_path, 'w', encoding='utf-8') as f:
        json.dump(simple_articles, f, ensure_ascii=False, indent=2)
    print(f"批次 {batch_id} 已保存 {len(simple_articles)} 篇文章到 {file_path}")


def parse_pubmed_article(element) -> Dict[str, Any]:
    """
    解析单个PubMed文章元素，提取批量处理所需的字段

    Args:
        element: XML元素

    Returns:
        Dict: 包含PMID, JournalTitle, ArticleTitle, AbstractText, Year的字典
    """
    data = {}

    # 解析MedlineCitation部分
    medline_citation = element.find('MedlineCitation')
    if medline_citation is not None:
        # 提取PMID
        pmid_elem = medline_citation.find('PMID')
        if pmid_elem is not None:
            data['PMID'] = pmid_elem.text

        # 提取年份
        date_completed = medline_citation.find('DateCompleted')
        if date_completed is not None:
            year_elem = date_completed.find('Year')
            if year_elem is not None:
                data['Year'] = year_elem.text

        # 提取文章信息
        article = medline_citation.find('Article')
        if article is not None:
            # 提取期刊标题
            journal = article.find('Journal')
            if journal is not None:
                journal_title = journal.find('Title')
                if journal_title is not None:
                    data['JournalTitle'] = journal_title.text

            # 提取文章标题
            title_elem = article.find('ArticleTitle')
            if title_elem is not None:
                data['ArticleTitle'] = title_elem.text

            # 提取摘要
            abstract = article.find('Abstract')
            if abstract is not None:
                abstract_text = abstract.find('AbstractText')
                if abstract_text is not None:
                    data['AbstractText'] = abstract_text.text

    return data

